Skip to content

Fix DoRA extraction crash on 1-D weights (x_pad_token, norm weights)#16

Open
Booyaka101 wants to merge 1 commit into
silveroxides:mainfrom
Booyaka101:fix/dora-extract-1d-tensor-crash
Open

Fix DoRA extraction crash on 1-D weights (x_pad_token, norm weights)#16
Booyaka101 wants to merge 1 commit into
silveroxides:mainfrom
Booyaka101:fix/dora-extract-1d-tensor-crash

Conversation

@Booyaka101

Copy link
Copy Markdown

Fixes #15.

Extracting a DoRA from models that contain 1-D weights (e.g. z-image's x_pad_token, or LayerNorm/RMSNorm .weight) crashes with:

Dimension out of range (expected to be in range of [-1, 0], but got 1)
  nodes/dora_extract_wd.py -> torch.linalg.norm(tensor_a, dim=1, keepdim=True)

In the matched-shape branch, is_conv_dim = tensor_a.ndim == 4 sends every non-4-D tensor to torch.linalg.norm(..., dim=1). A 1-D tensor has no dim=1, so it raises and aborts the whole extraction. The function already discards 1-D tensors a few lines later (the weight_diff.ndim < 2 skip), so this just bails out before the norm instead.

Verified with a CPU repro of the exact branch:

1-D input (x_pad_token / norm.weight):  before -> IndexError (dim 1) ;  after -> skipped
2-D input:                              still processed
4-D conv input:                         still processed

The matched-shape branch computed torch.linalg.norm(tensor_a, dim=1)
for every non-4-D tensor, including 1-D weights (e.g. x_pad_token,
norm .weight). A 1-D tensor has no dim=1, so this raised
'Dimension out of range ... but got 1' and aborted the whole
extraction. The code already skips 1-D tensors a few lines later, so
bail out before the norm instead.

Fixes silveroxides#15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Extraction fails on models with single dimension layers.

1 participant